package com.k2data.qa.zap;

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.logging.Logger;

import org.jsoup.Jsoup;
import org.jsoup.nodes.Document;
import org.jsoup.nodes.Element;
import org.jsoup.select.Elements;
import org.zaproxy.clientapi.core.ApiResponse;
import org.zaproxy.clientapi.core.ApiResponseElement;
import org.zaproxy.clientapi.core.ApiResponseList;
import org.zaproxy.clientapi.core.ApiResponseSet;
import org.zaproxy.clientapi.core.ClientApi;
import org.zaproxy.clientapi.core.ClientApiException;

/**
 * An example of how to set up authentication via the API and get information about existing
 * configuration.
 *
 * <p>Some important aspects regarding the Authentication API:
 *
 * <ul>
 *   <li>since the AuthenticationMethods are loaded dynamically, there's no way to generate a
 *       'static' API for each auth method. That's why, when setting up the authentication method,
 *       depending on the method, different values are passed to the setAuthenticationMethod . This
 *       is where the getSupportedAuthenticationMethods and getAuthenticationMethodConfigParams
 *       methods come into play. Basically the first one gives a list of available/loaded
 *       authentication methods while the second one gives info about the parameters required to
 *       configure each authentication method type.
 *   <li>when setting up the authentication method for a context, the setAuthenticationMethod method
 *       is used. It takes the context id on which we're working, the name of the authentication
 *       method and a 'authMethodConfigParams' parameter which contains all the configuration for
 *       the method. The format of the value passed for 'authMethodConfigParams' matches the
 *       www-form-urlencoded style: parameterName = urlEncodedValue. Check out the referenced
 *       example to see how the configuration is build for the KMX login. The pseudocode
 *       for generating the config would be:<br>
 *       <code>
 *       paramA + "=" + urlEncode(paramAValue) + "&" + paramB + "=" + urlEncode(paramBValue) + ...
 *       </code>
 *   <li>for formBasedAuthentication, the places filled in with the credentials are marked via
 *       {%username%} and {%password%}, in either the requestUrl or the requestBody
 * </ul>
 */
public class FormBasedAuthentication {
	private static final Logger logger = Logger.getLogger(FormBasedAuthentication.class.getName());

	private static final String ZAP_ADDRESS = "192.168.64.134";
    private static final int ZAP_PORT = 8088;
    private static final String ZAP_API_KEY = null;
    
    private static final String TARGET = "http://192.168.130.249:5010/";
    private static final String URL_LOGIN = "https://192.168.130.249:8443/cas/login";
    private static final String API_GATEWAY = "192.168.130.249";
    
    private static final String USERNAME = "tadmin_default_tenant";
    private static final String PASSWORD = "123456";

    /**
     * list auth information
     * @param clientApi
     * @throws ClientApiException
     */
    private static void listAuthInformation(ClientApi clientApi) throws ClientApiException {
        // Check out which authentication methods are supported by the API
        List<String> supportedMethodNames = new ArrayList<>();
        ApiResponseList authMethodsList = (ApiResponseList) clientApi.authentication.getSupportedAuthenticationMethods();
        for (ApiResponse authMethod : authMethodsList.getItems()) {
            supportedMethodNames.add(((ApiResponseElement) authMethod).getValue());
        }
        logger.info("Supported authentication methods: " + supportedMethodNames);

        // Check out which are the config parameters of the authentication methods
        for (String methodName : supportedMethodNames) {

            ApiResponseList configParamsList = (ApiResponseList) clientApi.authentication.getAuthenticationMethodConfigParams(methodName);

            for (ApiResponse r : configParamsList.getItems()) {
                ApiResponseSet set = (ApiResponseSet) r;
                logger.info("'"
                                + methodName
                                + "' config param: "
                                + set.getValue("name")
                                + " ("
                                + (set.getStringValue("mandatory").equals("true") ? "mandatory" : "optional")
                                + ")");
            }
        }
    }
    
    /**
     * create context
     * @param clientApi
     * @param contextName
     * @return
     */
    private static String createContext(ClientApi clientApi, String contextName) {
    	try {
			return extractontextId(clientApi.context.newContext(contextName));
		} catch (ClientApiException e) {
			e.printStackTrace();
		}
    	return null;
    }

    /**
     * list user config information
     * @param clientApi
     * @param contextId
     * @throws ClientApiException
     */
    private static void listUserConfigInformation(ClientApi clientApi, String contextId) throws ClientApiException {
        // Check out which are the config parameters required to set up an user with the currently
        // set authentication methods
        ApiResponseList configParamsList =
                (ApiResponseList)
                        clientApi.users.getAuthenticationCredentialsConfigParams(contextId);

        StringBuilder sb = new StringBuilder("Users' config params: ");
        for (ApiResponse r : configParamsList.getItems()) {
            ApiResponseSet set = (ApiResponseSet) r;
            sb.append(set.getValue("name")).append(" (");
            sb.append((set.getStringValue("mandatory").equals("true") ? "mandatory" : "optional"));
            sb.append("), ");
        }
        logger.info(sb.deleteCharAt(sb.length() - 2).toString());
    }

    /**
     * set logged in indicator
     * @param clientApi
     * @param contextId
     * @throws ClientApiException
     */
    private static void setLoggedInIndicator(ClientApi clientApi, String contextId) throws ClientApiException {
        // Prepare values to set, with the logged in indicator as a regex matching the logout link
        String loggedInIndicator = "Log In Successful";

        // Actually set the logged in indicator
        clientApi.authentication.setLoggedInIndicator(contextId, ".*" + java.util.regex.Pattern.quote(loggedInIndicator) + ".*" );

        // Check out the logged in indicator that is set
        logger.info( "Configured logged in indicator regex: " + 
        		((ApiResponseElement) clientApi.authentication.getLoggedInIndicator(contextId)).getValue());
    }
    
    /**
     * set logged out indicatior
     * @param clientApi
     * @param contextId
     * @throws ClientApiException
     */
    private static void setLoggedOutIndicator(ClientApi clientApi, String contextId) throws ClientApiException {
        // Prepare values to set, with the logged in indicator as a regex matching the logout link
        String loggedOutIndicator = URL_LOGIN;

        // Actually set the logged out indicator
        clientApi.authentication.setLoggedOutIndicator(contextId, ".*" + java.util.regex.Pattern.quote(loggedOutIndicator) + ".*" );

        // Check out the logged out indicator that is set
        logger.info( "Configured logged out indicator regex: " + 
        		((ApiResponseElement) clientApi.authentication.getLoggedOutIndicator(contextId)).getValue());
    }

    /**
     * set FormBasedAuthentication for KMX
     * @param clientApi
     * @param contextId
     * @throws ClientApiException
     * @throws UnsupportedEncodingException
     */
    private static void setFormBasedAuthenticationForKMX(ClientApi clientApi, String contextId)
            throws ClientApiException, UnsupportedEncodingException {
        // Setup the authentication method
        String loginUrl = URL_LOGIN;

        // token
        String token = extractLoginTokenForCAS(loginUrl);
        
        String loginRequestData = "username={%username%}&password={%password%}&execution=" + token + "&_eventId=submit";

        // Prepare the configuration in a format similar to how URL parameters are formed. This
        // means that any value we add for the configuration values has to be URL encoded.
        StringBuilder formBasedConfig = new StringBuilder();
        formBasedConfig.append("loginUrl=").append(URLEncoder.encode(loginUrl, "UTF-8"));
        formBasedConfig
                .append("&loginRequestData=")
                .append(URLEncoder.encode(loginRequestData, "UTF-8"));

        logger.info(
                "Setting form based authentication configuration as: "
                        + formBasedConfig.toString());
        clientApi.authentication.setAuthenticationMethod(
                contextId, "formBasedAuthentication", formBasedConfig.toString());

        // Check if everything is set up ok
        logger.info(
                "Authentication config: "
                        + clientApi.authentication.getAuthenticationMethod(contextId).toString(0));
    }

    /**
     * set user auth config for KMX
     * @param clientApi
     * @param contextId
     * @return
     * @throws ClientApiException
     * @throws UnsupportedEncodingException
     */
    private static String setUserAuthConfigForKMX(ClientApi clientApi, String contextId)
            throws ClientApiException, UnsupportedEncodingException {
        // Prepare info
        String user = "tadmin";
        String username = USERNAME;
        String password = PASSWORD;
        
        // Make sure we have at least one user
        String userId = extractUserId(clientApi.users.newUser(contextId, user));
        clientApi.users.setUserEnabled(contextId, userId, "True");
        
        // Prepare the configuration in a format similar to how URL parameters are formed. This
        // means that any value we add for the configuration values has to be URL encoded.
        StringBuilder userAuthConfig = new StringBuilder();
        userAuthConfig.append("username=").append(URLEncoder.encode(username, "UTF-8"));
        userAuthConfig.append("&password=").append(URLEncoder.encode(password, "UTF-8"));

        logger.info("Setting user authentication configuration as: " + userAuthConfig.toString());
        clientApi.users.setAuthenticationCredentials(contextId, userId, userAuthConfig.toString());

        // Check if everything is set up ok
        logger.info("Authentication config: " + clientApi.users.getUserById(contextId, userId).toString(0));
        
        return userId;
    }

    /**
     * extract userId
     * @param response
     * @return
     */
    private static String extractUserId(ApiResponse response) {
        return ((ApiResponseElement) response).getValue();
    }
    
    /**
     *get user id
     * @param response the ZAP API response code
     * @return the user ID of the context
     *
     */
    private static String extractontextId(ApiResponse response) {
        return ((ApiResponseElement) response).getValue();
    }
    
    /**
     * include sites in context
     * @param api
     */
    private static void setIncludeInContext(ClientApi api, String contextname) {
    	String httpSites = "http://" + API_GATEWAY + ".*";
    	String httpsSites = "https://" + API_GATEWAY + ".*";
    	
    	try {
			api.context.includeInContext(contextname, httpSites);			
			api.context.includeInContext(contextname, httpsSites);
			logger.info("IncludeInContext: " + httpSites + "\t" + httpsSites);
		} catch (ClientApiException e) {
			e.printStackTrace();
		}
    }
    
    /**
     * extract login token from html
     * @param url
     * @return
     */
	private static String extractLoginTokenForCAS(String url) {
		String token = null;
		
		try {
			Document doc = Jsoup.connect(url).validateTLSCertificates(false).get();
			Elements elements = doc.getElementsByAttributeValue("name", "execution");
			if(elements.size() == 1) {
				Element element = elements.first();
				token = element.attr("value");
				logger.info("execution: " + token);
			}
		} catch (IOException e1) {
			e1.printStackTrace();
		}
		
		return token;
	}
	
	/**
	 * active scan urls
	 * @param api
	 * @param contextId
	 * @param userId
	 * @param urls
	 */
	private static void scanUrls(ClientApi api, String contextId, String userId, List<String> urls) {
		
		if(urls != null) {
			for(String scanUrl : urls) {
				try {
					logger.info("ative scan url" + scanUrl);
					
					// access url before active scan
					api.accessUrl(scanUrl);
					
					// active scan
					ApiResponse resp = api.ascan.scanAsUser(scanUrl, contextId, userId, "True", "Default Policy", null, null);
					
					// The scan now returns a scan id to support concurrent scanning
	                String scanid = ((ApiResponseElement) resp).getValue();
	                int progress;
	 
	                // Poll the status until it completes
	                while (true) {
	                    Thread.sleep(5000);
	                    progress = Integer.parseInt(  ((ApiResponseElement) api.ascan.status(scanid)).getValue());
	                    logger.info("Active Scan progress : " + progress + "%");
	                    if (progress >= 100) {
	                        break;
	                    }
	                }
	                logger.info("Active Scan complete");
	 
	                logger.info("Alerts:");
	                logger.info(new String(api.core.xmlreport(), StandardCharsets.UTF_8));
				} catch (ClientApiException | InterruptedException e) {
					e.printStackTrace();
				}
			}
		}
	}
	
	/**
	 * import urls and scan
	 * @param api
	 * @param contextId
	 * @param userId
	 * @param urlFilePath
	 */
	private static void importAndScan(ClientApi api, String contextId, String userId, String urlFilePath){
		List<String> urlList = importUrls(api, urlFilePath);
		scanUrls(api, contextId, userId, urlList);
	}
	
    
    /**
     * scan target
     * @param api
     */
    private static void spiderAndScan(ClientApi api, String contextId, String userId) {
    	try {
    		String url = TARGET + "config.js";
    		
            // Start spidering the target
            logger.info("Spider : " + url);
            
            // It's not necessary to pass the ZAP API key again, already set when creating the ClientApi.
            ApiResponse resp =  api.spider.scanAsUser(contextId, userId, url, null, null, null);
            String scanid;
            int progress;
            
            // The scan now returns a scan id to support concurrent scanning
            scanid = ((ApiResponseElement) resp).getValue();
            
            // Poll the status until it completes
            while (true) {
                Thread.sleep(5000);
                progress =
                        Integer.parseInt(
                                ((ApiResponseElement) api.spider.status(scanid)).getValue());
                logger.info("Spider progress : " + progress + "%");
                if (progress >= 100) {
                    break;
                }
            }
            logger.info("Spider complete");
            
            // Give the passive scanner a chance to complete
            Thread.sleep(2000);

            // get all urls
            resp = api.spider.allUrls();
            List<ApiResponse> responseList = ((ApiResponseList) resp).getItems();
            List<String> urlList = new ArrayList<String>();
            for(ApiResponse response : responseList) {
            	String scanUrl = ((ApiResponseElement) response).getValue();
            	urlList.add(scanUrl);            	 
            }
            
            // scan urls
            scanUrls(api, contextId, userId, urlList);
        } catch (Exception e) {
            logger.info("Exception : " + e.getMessage());
            e.printStackTrace();
        }
    }
    
    /**
     * import urls from file
     * @param api
     * @param filepath
     */
    private static List<String> importUrls(ClientApi api, String filepath) {
    	List<String> urlList = new ArrayList<>();  
    	try {
//    		ApiResponse resp = api.importurls.importurls(filepath);
//    		logger.info("import urls " + ((ApiResponseElement) resp).getValue());
    		   		 		
    		String readline = null;
    		
    		BufferedReader br = new BufferedReader(new FileReader(filepath));
    		while((readline = br.readLine()) != null) {
    			urlList.add(readline);
    		}		
    		
    		br.close();
		} catch (IOException e) {
			e.printStackTrace();
		}
    	
    	return urlList;
    }

    /**
     * The main method.
     *
     * @param args the arguments
     * @throws Exception if an error occurred while accessing the API
     */
    public static void main(String[] args) throws Exception {	
    	String urlFilePath = "/home/p/urls.txt";
    	
        ClientApi clientApi = new ClientApi(ZAP_ADDRESS, ZAP_PORT, ZAP_API_KEY);
        String contextName = "KMX_Context_" + System.currentTimeMillis();
        
        String contextId = createContext(clientApi, contextName);
        listAuthInformation(clientApi);
        logger.info("------------------------------------------------------------------------------");
        setFormBasedAuthenticationForKMX(clientApi, contextId);
        logger.info("------------------------------------------------------------------------------");
        setLoggedInIndicator(clientApi, contextId);
        logger.info("------------------------------------------------------------------------------");
        setLoggedOutIndicator(clientApi, contextId);
        logger.info("------------------------------------------------------------------------------");
        listUserConfigInformation(clientApi, contextId);
        logger.info("------------------------------------------------------------------------------");
        String userId = setUserAuthConfigForKMX(clientApi, contextId);
        logger.info("------------------------------------------------------------------------------");       
        setIncludeInContext(clientApi, contextName);
        logger.info("------------------------------------------------------------------------------");
        spiderAndScan(clientApi, contextId, userId);
        logger.info("------------------------------------------------------------------------------");
        importAndScan(clientApi, contextId, userId, urlFilePath);
    }
}
